#!/usr/bin/env python

import rospy
import tf
import numpy as np
from geometry_msgs.msg import PoseArray,PoseStamped,Pose,PoseWithCovarianceStamped,TransformStamped
from ros_numpy import numpify,msgify
from nav_msgs.msg import Path
from tf.transformations import quaternion_from_matrix
import math
from sensor_msgs.msg import PointCloud2,PointCloud
import tf2_ros
from numpy.linalg import inv
from utils import calculate_angle1


def calculate_angle(vector1, vector2):
    dot_product = sum(x * y for x, y in zip(vector1, vector2))
    magnitude1 = math.sqrt(sum(x**2 for x in vector1))
    magnitude2 = math.sqrt(sum(x**2 for x in vector2))
    
    cosine_angle = dot_product / (magnitude1 * magnitude2)
    angle_in_radians = math.acos(cosine_angle)
    angle_in_degrees = math.degrees(angle_in_radians)
    
    return angle_in_radians

def np_to_pc2_msg(pc_np):
    data=np.zeros(pc_np.shape[1],
    dtype=[('x',np.float32), ('y',np.float32), ('z',np.float32)])
    data['x']=pc_np[0,]
    data['y']=pc_np[1,]
    data['z']=pc_np[2,]

    pc_msg=msgify(PointCloud2,data)
    pc_msg.is_dense=False

    return pc_msg


def calculate_obs_cov(T,s_x,s_yz,d_min):
    R=T[:3,:3]
    t=T[:3,3]
    d=np.sqrt(np.sum(t**2))
    d_min=0.3
    V=np.zeros([3,3])
    V[0,0]=(d-d_min)**2*s_x
    V[1,1]=(d-d_min)**2*s_yz
    V[2,2]=(d-d_min)**2*s_yz
    
    S= R @ V @R.T
    
    return S

def kalman_filter(P,Q,S,F=np.eye(3)):
    
    P_next=inv(inv(F@P@F.T+Q)+inv(S))

    return P_next

class TFPublisher:
    def __init__(self):
        # Initialize the ROS node
        rospy.init_node('tf_publisher')
        self.t0=rospy.get_time()
        
        self.tf_buffer = tf2_ros.Buffer()
        self.tf_listener = tf2_ros.TransformListener(self.tf_buffer)

        # Create a TF broadcaster
        self.broadcaster = tf.TransformBroadcaster()

        # Set the publishing rate
        self.rate = rospy.Rate(10)  # 10 Hz
        
        self.body_cmd_pose=np.eye(4)
        self.pc_W=None
        self.pc_B=None
        self.body_cmd_pose[:3,3]=np.array([-0.0,0.01,1.323])
        self.x=0
        rospy.Subscriber('/slot/view_array',PoseArray,self.cmd_pose_arr_cb)
        self.pc_W_pub=rospy.Publisher('/slot/pointcloud2_vrpn',PointCloud2,queue_size=10)
    
        self.pc_B_pub=rospy.Publisher('/slot/pointcloud2_body',PointCloud2,queue_size=10)
        self.pc_W_vis_pub=rospy.Publisher('/slot/pointcloud2_visible',PointCloud2,queue_size=10)
        self.pc_count_list=[]
        
        self.body_path=Path()
        self.target_path=Path()
        self.body_path_pub=rospy.Publisher('/slot/body/path',Path,queue_size=10)
        self.target_path_pub=rospy.Publisher('/slot/target/path',Path,queue_size=10)
        self.body_path.header.frame_id='vrpn'
        self.target_path.header.frame_id='vrpn'
        
        self.kf_Q=np.eye(3)*0.1
        self.kf_V=np.array([100.,1.,1.])
        self.kf_P=np.eye(3)
        
    
    def pc_generator(self):
        xc=16.0
        zc=3.3
        yc=-6.7
        density_l=0.3
        density_r=2.3
    
        
        z_l=np.arange(-6.7,0,density_l)
        z_r=np.arange(0,+6.7,density_r)
        z=np.hstack([z_l,z_r])+zc
        
        y_l=np.arange(-16.7,0,density_l)
        y_r=np.arange(0,+16.7,density_r)
        y=np.hstack([y_l,y_r])+yc
        
        pc_W=np.zeros([3,z.shape[0]*y.shape[0]])
        pc_W[0,:]=xc
        
        i=0
        for zi in z:
            for yi in y:
                pc_W[1,i]=yi
                pc_W[2,i]=zi
                
                i=i+1

        pc_W_msg=np_to_pc2_msg(pc_W)
        pc_W_msg.header.frame_id='vrpn'
        self.pc_W=pc_W
        
        self.pc_W_pub.publish(pc_W_msg)
    #############################################
    #publish visable pc in body frame
        pc_W_aug=np.vstack([self.pc_W,np.ones(self.pc_W.shape[1])])
        pc_B_aug=np.linalg.inv(self.body_cmd_pose) @ pc_W_aug
        pc_B=pc_B_aug[:3,:]
        self.pc_B=pc_B
    
        
        
    def cmd_pose_arr_cb(self,pose_arr_msg):
        
        #msg=Path()
        cmd_pose_msg=pose_arr_msg.poses[1]
        body_pose_stamped=PoseStamped()
        body_pose_stamped.header.frame_id='vrpn'
        body_pose_stamped.pose=cmd_pose_msg
        self.body_path.poses.append(body_pose_stamped)
        self.body_cmd_pose=numpify(cmd_pose_msg)

    def publish_transforms(self):
        while not rospy.is_shutdown():
            
            self.pc_generator()
            # Publish the transformation from 'vrpn' to 'target'
            
            body_cmd_pos=self.body_cmd_pose[:3,3]
            body_cmd_qua=quaternion_from_matrix(self.body_cmd_pose)
            
            T_T_V=np.eye(4)
            T_T_V[:3,3]=np.array([self.x+2.5, -0.01, 1.3])
            
            r=2.0
            a=2.0
            b=4.0
            omega=0.13
            cx=0.
            cy=0.
            target_pose=Pose()
            target_pose.position.y=a*math.cos(omega*(rospy.get_time()-self.t0))+cy
            target_pose.position.x=b*math.sin(omega*(rospy.get_time()-self.t0))+cx
            target_pose.position.z=1.3
            
            
            #T_T_V=numpify(target_pose) # circle
            
            target_pose=msgify(Pose,T_T_V)
            
            target_pose_stamped=PoseStamped()
            target_pose_stamped.header.frame_id='vrpn'
            target_pose_stamped.pose=target_pose
            self.target_path.poses.append(target_pose_stamped)
            
            self.body_path_pub.publish(self.body_path)
            self.target_path_pub.publish(self.target_path)
            
            
            self.broadcaster.sendTransform(
                T_T_V[:3,3],           # Translation
                (0.0, 0.0, 0.0, 1.0),      # Rotation
                rospy.Time.now(),          # Time stamp
                'target',                  # Child frame ID
                'vrpn'                     # Parent frame ID
            )

            # Publish the transformation from 'target' to 'body'
            self.broadcaster.sendTransform(
                body_cmd_pos,           # Translation
                body_cmd_qua,      # Rotation
                rospy.Time.now(),          # Time stamp
                'body',                    # Child frame ID
                'vrpn'                   # Parent frame ID
            )

            self.x+=0.03 #0.3 m/s
            
            ######
            
                    #check visibility
                    
            T_B_V=self.body_cmd_pose
            T_T_B=np.linalg.inv(T_B_V) @ T_T_V
            
            T_B_vector=(T_T_B @ np.array([0,0,0,1]))[:3]
            
            d_T_B=np.sqrt(np.sum(T_B_vector**2))
            
            pc_B_visable_list=[]
            pc_W_visable_list=[]
            pc_B=self.pc_B
            pc_W=self.pc_W
            for col in range(pc_B.shape[1]):
                x=pc_B[0,col]
                y=pc_B[1,col]
                z=pc_B[2,col]
                if np.abs(calculate_angle(pc_B[:,col],np.array([1.,0.,0.])))<np.pi/4:
                    if np.abs(calculate_angle(pc_B[:,col],T_B_vector))>0.2313/d_T_B:
                        pc_B_visable_list.append(pc_B[:,col])
                        pc_W_visable_list.append(pc_W[:,col])
                        
            
            if(len(pc_B_visable_list)>0):
            
                pc_B=np.vstack(pc_B_visable_list).T
                pc_W_visable=np.vstack(pc_W_visable_list).T
            
            pc_B_msg=np_to_pc2_msg(pc_B)
            pc_B_msg.header.stamp=rospy.Time.now()
            pc_B_msg.header.frame_id="body"
        
            self.pc_B_pub.publish(pc_B_msg)
            
            ##################################
            pc_W_visable_msg=np_to_pc2_msg(pc_W_visable)
            pc_W_visable_msg.header.stamp=rospy.Time.now()
            pc_W_visable_msg.header.frame_id="vrpn"
            self.pc_W_vis_pub.publish(pc_W_visable_msg)
            
            ####################
            self.pc_count_list.append(pc_B.shape[1])
            pc_count_arr=np.array(self.pc_count_list)
            
            #np.save('pc_count_3.npy',pc_count_arr)
            
            self.rate.sleep()

if __name__ == '__main__':
    try:
        tf_publisher = TFPublisher()
        tf_publisher.publish_transforms()
    except rospy.ROSInterruptException:
        pass